#-----------------------------------------------------#
#       评价框架：计算mIoU
#-----------------------------------------------------#
import tensorflow as tf
import tensorflow.keras as keras

import PIL.Image as Image
import os
import time
import shutil
from tensorflow.python.ops.gen_io_ops import save
from tqdm import tqdm
import numpy as np

from net_main import XX_NET as XX_NET_proto
import utils.utils_evaluate as utils_evaluate

'''
    # 评价框架实现：
        1.将测试集的样本计算得到分割图
        2.计算分割图与标签的IoU
'''
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)


#------- 设置面板：全局变量 ------------------------------------#
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)          # 屏蔽tensorflow的warning

# tf.config.threading.set_inter_op_parallelism_threads(2) # 控制计算使用的线程 i.e. CPU资源
# tf.config.threading.set_intra_op_parallelism_threads(2)

# 1.准备模型
class XX_NET(XX_NET_proto):
    def __init__(self):
        super()._defaults['num_classes']         = 1                    # 修改类属性：训练是否需要加载分类名称表
        super()._defaults['load_classname_list'] = True                 # 修改是否加载分类名字清单。设置为True后，self.num_classes预设值失效，初始化时自动根据classes_path文件更新
        super()._defaults['load_weights']        = True                # 修改类属性：训练是否需要加载权重
        super()._defaults['model_weights_path']  = './logs/w_e.h5'
        super().__init__()                                              # 实例化方法
    def detect_image(self, image):
        ###----- 1.前处理
        resized_image = image.resize(self.net_input_shape[0:2], Image.BICUBIC)

        photo = np.asarray(resized_image)
        photo = np.expand_dims(photo, axis=0)
        photo = self.net_preprocess_input(photo)

        ###----- 2.网络
        pred = self.get_net_pred(photo) # (None, 480, 480, 21)
        pred = pred[0].numpy()

        cond_1 = pred < 0
        cond_2 = pred > 1
        assert not(cond_1.any() or cond_2.any()), 'There is illegal value in pred, which is not in range of [0,1]'
        
        ###----- 3.后处理
        pred_classIndex = pred.argmax(axis=-1) # (480, 480)
        pred_classIndex = pred_classIndex + 1  # 直接的预测结果中背景类被标记为第20类，需要调整为第0类
        pred_classIndex[pred_classIndex == 21] = 0

        seg_image = np.uint8(pred_classIndex)
        seg_image = Image.fromarray(seg_image).resize((image.width, image.height), Image.NEAREST)

        return seg_image
        
xx_net = XX_NET()
net_model = xx_net.net_model
detect_image = xx_net.detect_image
# net_model.summary()

xx_net.class_names.insert(0, 'background')

# 2.准备数据
annotation_path_test   = './model_data/test.txt'

dataset_image_dir = 'D:/【AI】/Datasets/CV_ds/VOC2012/VOCtrainval_11-May-2012/VOCdevkit/VOC2012/JPEGImages/'
dataset_label_dir = 'D:/【AI】/Datasets/CV_ds/VOC2012/VOCtrainval_11-May-2012/VOCdevkit/VOC2012/SegmentationClass/'

evaluate_time = time.localtime()
evaluate_results_dir = './logs/evaluate_results_{}-{:0>2d}-{:0>2d}_{:0>2d}{:2>2d}/'.format(evaluate_time.tm_year, 
                                                                          evaluate_time.tm_mon,
                                                                          evaluate_time.tm_mday, 
                                                                          evaluate_time.tm_hour,
                                                                          evaluate_time.tm_min)
if os.path.exists(evaluate_results_dir):
    print('dir "{}" exists!'.format(evaluate_results_dir))
else:
    os.mkdir(evaluate_results_dir)
save_seg_image_dir = evaluate_results_dir + 'seg_images/'

cmp_labels = [0, 0, 0,          # 生成VOC标签格式的索引图需要用到的标记颜色的矩阵
              128, 0, 0,
              0, 128, 0,        
              128, 128, 0,      
              0, 0, 128,        
              128, 0, 128,      
              0, 128, 128,      
              128, 128, 128,    
              64, 0, 0,         
              192, 0, 0,        
              64, 128, 0,       
              192, 128, 0,      
              64, 0, 128,       
              192, 0, 128,      
              64, 128, 128,     
              192, 128, 128,    
              0, 64, 0,         
              128, 64, 0,       
              0, 192, 0,        
              128, 192, 0,      
              0, 64, 128]
#-------------------------------------------------------------#


###----- 1.测试集样本
with open(annotation_path_test) as f:
    test_samples = f.readlines()
test_samples = [c.strip() for c in test_samples]


###----- 2.生成预测分割图
if os.path.exists(save_seg_image_dir):
    shutil.rmtree(save_seg_image_dir)
    os.makedirs(save_seg_image_dir)
else:
    os.makedirs(save_seg_image_dir)

for img in tqdm(test_samples, desc='Generating predicted seg_images: '):
    try:
        image = Image.open(dataset_image_dir + img + '.jpg')
        image = image.convert('RGB')
    except:
        print(img + ': open error!')
        continue
    else:
        seg_image = detect_image(image)
        seg_image.putpalette(cmp_labels) # 将灰度图转为索引图，与VOC的标记图片格式一致
        seg_image.save(save_seg_image_dir + img + '.png')

print('Finished generating predicted seg_images of test_samples.')


###----- 3.计算mIoU
mIoU, reports = utils_evaluate.compute_mIoU(gt_dir        = dataset_label_dir, 
                                            pred_dir      = save_seg_image_dir, 
                                            png_name_list = test_samples, 
                                            num_classes   = xx_net.num_classes + 1, 
                                            name_classes  = xx_net.class_names)

with open(file=evaluate_results_dir + '_evaluate_report.txt', mode='x') as f:
    # for i in reports:
    #     f.writelines(i)
    pass

'''
    VOC标记使用的索引图，参考：
        https://blog.csdn.net/dcrmg/article/details/94457607
'''